import pandas as pd
from pandasai.dataframe.base import DataFrame
from pandasai.config import APIKeyManager,ConfigManager
from pandasai.agent.base import Agent
from extensions.llms.litellm.pandasai_litellm import LiteLLM
# step1: 配置api和llm
# 讯飞的有每日的免费额度； 火山的开启了安全体验，后期可以给它关了
key="bce65545-613a-4cae-b81c-d8f7cc445d7d"
llm = LiteLLM(
    model="volcengine/doubao-seed-1-6-250615", # 这里必须写llmlite支持的模型名称
    api_key=key,
    temperature=0.2
)
ConfigManager.set({"llm":llm})
# step1: 加载 Excel 文件为 DataFrame
def read_csv(filepath) -> DataFrame:
    data = pd.read_excel(filepath)
    # table = get_table_name_from_path(filepath)
    return DataFrame(data) #, _table_name=table
data_path=r'test_excels/gwd_indicator.xlsx'
df = read_csv(data_path) # 传入必须是DataFrame类
# step3: 然后传入 Agent并聊天：
agent = Agent([df])
query="你好，我想知道这个数据里每一列数据的基础信息，包括数据类型、数值范围、分布等"
agent.chat(query)